from PyQt5.QtWidgets import (QMainWindow, QLabel, QPushButton, QVBoxLayout, 
                            QWidget, QMessageBox, QFileDialog)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from core.image_selector import ImageSelector
from core.ocr_engine import OCREngine
from core.translator import Translator
from core.image_generator import ImageGenerator
from ui.progress_dialog import ProgressDialog
import os

class TranslationThread(QThread):
    progress_updated = pyqtSignal(int, str)
    finished = pyqtSignal(bool, str)
    
    def __init__(self, files: list, is_folder: bool, output_dir: str):
        super().__init__()
        self.files = files
        self.is_folder = is_folder
        self.output_dir = output_dir
        self.ocr = OCREngine()
        self.translator = Translator(service='google')
    
    def run(self):
        try:
            # 如果是文件夹，获取文件夹下所有图片文件
            if self.is_folder and len(self.files) == 1:
                folder_path = self.files[0]
                self.files = []
                for file_name in os.listdir(folder_path):
                    if file_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                        self.files.append(os.path.join(folder_path, file_name))
                if not self.files:
                    self.finished.emit(False, "文件夹中没有找到图片文件")
                    return
            
            total = len(self.files)
            for i, file_path in enumerate(self.files):
                # 更新进度
                self.progress_updated.emit(int((i+1)/total*100), f"处理中: {os.path.basename(file_path)}")
                
                # OCR识别
                ocr_results = self.ocr.detect_text(file_path)
                if not ocr_results:
                    continue
                
                # 翻译文本
                texts = [result["text"] for result in ocr_results]
                translations = self.translator.batch_translate(texts)
                
                # 确保输出目录存在
                os.makedirs(self.output_dir, exist_ok=True)
                
                # 生成图片
                output_path = os.path.join(
                    self.output_dir, 
                    f"translated_{os.path.basename(file_path)}"
                )
                ImageGenerator.generate_translated_image(
                    file_path, ocr_results, translations, output_path
                )
            
            self.finished.emit(True, f"处理完成！结果保存在: {self.output_dir}")
        except Exception as e:
            self.finished.emit(False, f"处理失败: {str(e)}")

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setup_ui()
        self.current_files = []
        self.is_folder = False
    
    def setup_ui(self):
        self.setWindowTitle("图片翻译系统")
        self.setGeometry(100, 100, 600, 400)
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        layout = QVBoxLayout()
        
        # 添加控件
        self.label = QLabel("请选择要翻译的图片或文件夹")
        self.label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.label)
        
        self.import_btn = QPushButton("导入图片")
        self.import_btn.clicked.connect(self.import_images)
        layout.addWidget(self.import_btn)
        
        self.translate_btn = QPushButton("开始翻译")
        self.translate_btn.clicked.connect(self.start_translation)
        self.translate_btn.setEnabled(False)
        layout.addWidget(self.translate_btn)
        
        central_widget.setLayout(layout)
    
    def import_images(self):
        self.current_files, self.is_folder = ImageSelector.select_images(self)
        if self.current_files:
            self.translate_btn.setEnabled(True)
            if self.is_folder:
                self.label.setText(f"已选择文件夹: {self.current_files[0]}")
            else:
                self.label.setText(f"已选择 {len(self.current_files)} 个图片文件")
    
    def start_translation(self):
        # 选择输出目录
        output_dir = QFileDialog.getExistingDirectory(self, "选择输出目录")
        if not output_dir:
            return
        
        # 创建并显示进度对话框
        self.progress_dialog = ProgressDialog(self)
        self.progress_dialog.show()
        
        # 创建并启动工作线程
        self.worker = TranslationThread(self.current_files, self.is_folder, output_dir)
        self.worker.progress_updated.connect(self.progress_dialog.update_progress)
        self.worker.finished.connect(self.on_translation_finished)
        self.worker.start()
    
    def on_translation_finished(self, success: bool, message: str):
        self.progress_dialog.close()
        QMessageBox.information(self, "完成" if success else "错误", message)
        if success:
            self.translate_btn.setEnabled(False)
            self.label.setText("翻译完成！可以导入新的图片")